import numpy as np
import cv2
import os

from typing import List, Tuple

from SimpleProcessor import Simple_Processor


class Simple_DataSet_Loader(object):
    def __init__(self, preprocessors: List[Simple_Processor] = None):
        """
        初始化数据集加载器

        :param preprocessors: 数据集预处理器
        """
        self.preprocessors = preprocessors if preprocessors else []

    def load(self, train_image_path: List[str]) -> Tuple[np.ndarray, np.ndarray]:
        """
        数据集加载器

        :param train_image_path: 字符串列表（可考虑改为字符串生成器）
        :return:
        """
        data = []
        labels = []
        for num, raw_image in enumerate(train_image_path, start=1):
            image = cv2.imread(raw_image)
            label = raw_image.split(os.path.sep)[-1].split('.')[0]
            if self.preprocessors:
                for p in self.preprocessors:
                    image = p.processor(image)
            data.append(image)
            labels.append(label)
            print(f"[info]:{num} / {len(train_image_path)} {label}'s image have loaded")
        return np.array(data), np.array(labels)


if __name__ == '__main__':
    from pathlib import Path
    s = Path('../dataset/dog_cat/test')
    ss = Simple_DataSet_Loader()
    print(ss.load([str(i) for i in s.iterdir()]))
